import argparse
import torch
from tensorboardX import SummaryWriter
from models.vmunet.vmunet import VMUNet

import os
import sys

from configs.config_setting_tdmaterials import setting_config

from trainer import trainer_tdmaterials
from utils import cal_params_flops, get_logger, log_config_info, set_seed

import warnings
warnings.filterwarnings("ignore")


parser = argparse.ArgumentParser()
parser.add_argument('--material', type=str,
                    default='Graphene', help='Type of 2D material')
parser.add_argument('--num_classes', type=int,
                    default=4, help='output channel of network')
parser.add_argument('--max_epochs', type=int,
                    default=300, help='maximum epoch number to train')
parser.add_argument('--batch_size', type=int,
                    default=8, help='batch_size per gpu')
parser.add_argument('--n_gpu', type=int,
                    default=1, help='total gpu')
parser.add_argument('--img_size', type=int,
                    default=512, help='input patch size of network input')
parser.add_argument('--seed', type=int,
                    default=1234, help='random seed')
parser.add_argument('--pretrained_path', type=str,
                    default='', help='if using pretrained, please enter the path of weights')
parser.add_argument('--cost', type=str,
                    default='ce-dice', help='which cost function to use')
args = parser.parse_args()

config = setting_config(args)


if __name__ == '__main__':
    
    
    print('#----------Creating logger----------#')
    sys.path.append(config.work_dir + '/')
    log_dir = os.path.join(config.work_dir, 'log')
    checkpoint_dir = os.path.join(config.work_dir, 'checkpoints')
    resume_model = os.path.join(checkpoint_dir, 'latest.pth')
    outputs = os.path.join(config.work_dir, 'outputs')
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)
    if not os.path.exists(outputs):
        os.makedirs(outputs)

    logger = get_logger('train', log_dir)
    writer = SummaryWriter(config.work_dir + 'summary')

    log_config_info(config, logger)


    print('#----------GPU init----------#')
    set_seed(config.seed)
    torch.cuda.empty_cache()

    print('#----------Prepareing Model----------#')
    model_cfg = config.model_config
    if config.network == 'vmunet':
        model = VMUNet(
            num_classes=model_cfg['num_classes'],
            input_channels=model_cfg['input_channels'],
            depths=model_cfg['depths'],
            depths_decoder=model_cfg['depths_decoder'],
            drop_path_rate=model_cfg['drop_path_rate'],
            load_ckpt_path=model_cfg['load_ckpt_path'],
        )
        model.load_from()
        
    else: raise Exception('network in not right!')

    model = model.cuda()

    cal_params_flops(model, args.img_size, logger)

    trainer = {'TDMaterials': trainer_tdmaterials}
    trainer[config.datasets_name](config, model, resume_model, checkpoint_dir, logger, writer)
